#include "ide.h"
#include "sync.h"
#include "io.h"
#include "stdio.h"
#include "stdio-kernel.h"
#include "interrupt.h"
#include "memory.h"
#include "debug.h"
#include "console.h"
#include "timer.h"
#include "string.h"
#include "list.h"

/* 定义硬盘各寄存器的端口号 */
#define reg_data(channel)	    (channel->port_base + 0)
#define reg_error(channel)	    (channel->port_base + 1)
#define reg_sect_cnt(channel)	(channel->port_base + 2)
#define reg_lba_l(channel)	    (channel->port_base + 3)
#define reg_lba_m(channel)	    (channel->port_base + 4)
#define reg_lba_h(channel)	    (channel->port_base + 5)
#define reg_dev(channel)	    (channel->port_base + 6)
#define reg_status(channel)     (channel->port_base + 7)
#define reg_cmd(channel)	    (reg_status(channel))
#define reg_alt_status(channel) (channel->port_base + 0x206)
#define reg_ctl(channel)	    reg_alt_status(channel)

/* reg_status寄存器的一些关键位 */
#define BIT_STAT_BSY	 0x80	      // 硬盘忙
#define BIT_STAT_DRDY	 0x40	      // 驱动器准备好	 
#define BIT_STAT_DRQ	 0x8	      // 数据传输准备好了

/* device寄存器的一些关键位 */
#define BIT_DEV_MBS	0xa0	    // 第7位和第5位固定为1
#define BIT_DEV_LBA	0x40
#define BIT_DEV_DEV	0x10

/* 一些硬盘操作的指令 */
#define CMD_IDENTIFY	   0xec	    // identify指令
#define CMD_READ_SECTOR	   0x20     // 读扇区指令
#define CMD_WRITE_SECTOR   0x30	    // 写扇区指令

/* 定义可读写的最大扇区数,调试用的 */
#define max_lba ((80*1024*1024/512) - 1)	// 只支持80MB硬盘

uint8_t channel_cnt; // 硬盘 / 2 就是通道数量
struct ide_channel channels[2]; // 有两个 IDE 通道

// 用于记录总扩展分区的起始 LBA，初始为 0，partition_scan 时以此为标记
int32_t ext_lba_base = 0;

uint8_t p_no = 0, l_no = 0; // 用来记录硬盘主分区和逻辑分区的下标

struct list partition_list; // 分区队列

// 构建一个 16 字节大小的结构体，用来存储分区表项
struct partition_table_entry {
    uint8_t bootable;   // 是否可引导
    uint8_t start_head; // 起始磁头号
    uint8_t start_sec;  // 起始扇区号
    uint8_t start_chs;  // 起始柱面号
    uint8_t fs_type;    // 分区类型
    uint8_t end_head;   // 结束磁头号
    uint8_t end_sec;    // 结束扇区号
    uint8_t end_chs;    // 结束柱面号
    uint32_t start_lba; // 本分区起始扇区的 LBA 地址
    uint32_t sec_cnt;   // 本分区的扇区数目
} __attribute__ ((packed)); // 保证此结构是 16 字节大小

// 引导扇区，mbr 或 ebr 所在的扇区
struct boot_sector {
    uint8_t other[446]; // 引导扇区的代码 这里只是用来占位的 为了方便获取到分区表
    struct partition_table_entry partition_table[4]; // 分区表中有 4 项，每项 16 字节，共 64 字节
    uint16_t signature; // 引导扇区的结束魔数(即标志)是 0x55, 0xaa, 小端序得到实际值为 0xaa55
} __attribute__ ((packed));

// 选择要读写的是哪个硬盘
static void select_disk(struct disk* hd) {
    uint8_t reg_device = BIT_DEV_MBS | BIT_DEV_LBA;
    if(hd -> dev_no == 1) // 主盘为 0，从盘为 1
        reg_device |= BIT_DEV_DEV;
    outb(reg_dev(hd -> my_channel), reg_device);
}

// 向硬盘控制器写入起始扇区地址及要读写的扇区数
static void select_sector(struct disk* hd, uint32_t lba, uint8_t sec_cnt) {
    ASSERT(lba <= max_lba);
    struct ide_channel* channel = hd -> my_channel;
    // 写入要读写的扇区数
    outb(reg_sect_cnt(channel), sec_cnt);
    // 写入 LBA 地址(即扇区号)
    outb(reg_lba_l(channel), lba);      // LBA 地址的低 8 位
    outb(reg_lba_m(channel), lba >> 8); // LBA 地址的低 8~15 位
    outb(reg_lba_h(channel), lba >> 16);// LBA 地址的高 16 ~ 23

    // device 的低 0~3 位存储 LBA 地址的高 24~27 位
    // 因为寄存器无法单独写入，因此需要把整个 device 寄存器的值重新写入
    outb(reg_dev(channel), BIT_DEV_MBS | BIT_DEV_LBA | (hd -> dev_no == 1 ? BIT_DEV_DEV : 0) | lba >> 24);
}

// 向通道 channel 发送命令
static void cmd_out(struct ide_channel* channel, uint8_t cmd) {
    // 因为向硬盘发送命令后，硬盘工作需要时间，硬盘工作完成后会给我们发送一个中断，我们需要等待这个中断
    // 因此，需要将 channel -> expecting_intr 置为 true 用于标识是否等待该通道上的硬盘中断
    channel -> expecting_intr = true;
    outb(reg_cmd(channel), cmd);
}

// 硬盘读入 sec_cnt 个扇区的数据到 buf
static void read_from_sector(struct disk* hd, void* buf, uint8_t sec_cnt) {
    uint32_t size_in_size;
    if(sec_cnt == 0) // 因为 sec_cnt 为 8 位, 256 个扇区转成二进制为 1 0000 0000, 会自动舍去最高位，从而保留 8 位，因此 0 表示要读写 256 个扇区
        size_in_size = 256 * 512;
    else
        size_in_size = sec_cnt * 512;
    insw(reg_data(hd -> my_channel), buf, size_in_size / 2);
}

// 将 buf 中 sec_cnt 个扇区的数据写入硬盘
static void write2sector(struct disk* hd, void* buf, uint8_t sec_cnt) {
    uint32_t size_in_byte;
    if(sec_cnt == 0)
        size_in_byte = 256 * 512;
    else
        size_in_byte = sec_cnt * 512;
    outsw(reg_data(hd -> my_channel), buf, size_in_byte / 2);
}

// 等待 30 秒
static bool busy_wait(struct disk* hd) {
    struct ide_channel* channel = hd -> my_channel;
    uint16_t time_list = 30 * 1000;
    while(time_list -= 10 >= 0) {
        if(!(inb(reg_status(channel)) & BIT_STAT_BSY)) 
            return (inb(reg_status(channel)) & BIT_STAT_DRQ);
        else
            mtime_sleep(10);
    }
    return false;
}

// 从硬盘读取 sec_cnt 个扇区到 buf
void ide_read(struct disk* hd, uint32_t lba, void* buf, uint32_t sec_cnt) {
    ASSERT(lba <= max_lba);
    ASSERT(sec_cnt > 0);
    lock_acquire(&hd -> my_channel -> lock);

    select_disk(hd); // 先选择要操作的硬盘

    uint32_t secs_op;       // 每次操作的扇区数（每次最多不可超过 256 个）
    uint32_t secs_done = 0; // 已完成的扇区数
    
    while(secs_done < sec_cnt) {
        if((secs_done + 256) <= sec_cnt) // 若本次的最大操作数还未达到 sec_cnt，则直接操作最大值 256
            secs_op = 256;
        else
            secs_op = sec_cnt - secs_done; // 若本次的最大操作数已经溢出，则求它们之间的差

        // 写入待读入的扇区数和起始扇区号
        select_sector(hd, lba + secs_done, secs_op);

        // 执行的命令写入到 reg_cmd 寄存器
        cmd_out(hd -> my_channel, CMD_READ_SECTOR); // 准备开始读数据

        // 当硬盘已经开始工作后才可以阻塞自己，执行完上面的指令后硬盘开始工作
        // 因此此时可以将自己阻塞，等待硬盘完成操作后通过中断处理程序唤醒自己
        sema_down(&hd -> my_channel -> disk_done);

        // 检查硬盘状态是否可读
        if(!busy_wait(hd)) {
            char error[64];
            sprintf(error, "%s read sector %d failed!!!!!!\n", hd -> name, lba);
            PANIC(error);
        }

        // 把数据从硬盘的缓冲区读出
        read_from_sector(hd, (void*)((uint32_t)buf + secs_done * 512), secs_op);

        secs_done += secs_op;
    }
    lock_release(&hd -> my_channel -> lock);
}

// 将 buf 中 sec_cnt 扇区数据写入数据
void ide_write(struct disk* hd, uint32_t lba, void* buf, uint32_t sec_cnt) {
    ASSERT(lba <= max_lba);
    ASSERT(sec_cnt > 0);
    lock_acquire(&hd -> my_channel -> lock);

    // 先选择要操作的硬盘
    select_disk(hd);

    uint32_t secs_op;
    uint32_t secs_done = 0;

    while(secs_done < sec_cnt) {
        if(secs_done + 256 <= sec_cnt) 
            secs_op = 256;
        else 
            secs_op = sec_cnt - secs_done;

        // 向硬盘写入起始扇区地址以及要读写的扇区数
        select_sector(hd, lba + secs_done, secs_op);

        // 向硬盘发送命令
        cmd_out(hd -> my_channel, CMD_WRITE_SECTOR);

        // 检查硬盘状态是否可读
        if(!busy_wait(hd)) {
            char error[64];
            sprintf(error, "%s write sector %d failed!!!!!!\n", hd->name, lba);
            PANIC(error);
        }

        // 将数据写入硬盘
        write2sector(hd, (void*) ((uint32_t)buf + secs_done * 512), secs_op);

        // 在硬盘响应期间阻塞自己
        sema_down(&hd -> my_channel -> disk_done);

        secs_done += secs_op;
    }

    lock_release(&hd -> my_channel -> lock);
}

// 将 dst 中 len 个相邻字节交换位置后存入 buf
static void swap_pairs_bytes(const char* dst, char* buf, uint32_t len) {
    uint8_t idx;
    for(idx = 0; idx < len; idx += 2) {
        // buf 中存储 dst 中两两相邻元素交换位置后的字符串
        buf[idx + 1] = *dst++;
        buf[idx] = *dst++;
    }
    buf[idx] = '\0';
}

// 获得硬盘参数信息
static void identify_disk(struct disk* hd) {
    char id_info[512];
    select_disk(hd);
    cmd_out(hd -> my_channel, CMD_IDENTIFY);

    // 上面向硬盘发送完命令后，此时需要将自己阻塞，等待来着硬盘的中断
    sema_down(&hd -> my_channel -> disk_done);

    // 醒来后开始执行下面代码
    if(!busy_wait(hd)) {
        char error[64];
        sprintf(error, "%s identify failed!!!!!!\n", hd->name);
        PANIC(error);
    }

    // 读取一个扇区
    read_from_sector(hd, id_info, 1);

    char buf[64];
    uint8_t sn_start = 10 * 2, sn_len = 20, md_start = 27 * 2, md_len = 40;

    swap_pairs_bytes(&id_info[sn_start], buf, sn_len);
    printk("   disk %s info:\n      SN: %s\n", hd->name, buf);
    memset(buf, 0, sizeof(buf));

    swap_pairs_bytes(&id_info[md_start], buf, md_len);
    printk("      MODULE: %s\n", buf);
    uint32_t sectors = *(uint32_t*) &id_info[60 * 2];
    printk("      SECTORS: %d\n", sectors);
    printk("      CAPACITY: %dMB\n", sectors * 512 / 1024 / 1024);
}

// 扫描键盘 hd 中地址为 ext_lba 的扇区中的所有分区
static void partition_scan(struct disk* hd, uint32_t ext_lba) {
    struct boot_sector* bs = sys_malloc(sizeof(struct boot_sector));
    ide_read(hd, ext_lba, bs, 1); // 在硬盘 hd 起始扇区 LBA 地址为 ext_lba 的地方起，读取一个扇区，并且保存到 bs
    uint8_t part_idx = 0; // 当前分区索引
    struct partition_table_entry* p = bs -> partition_table;

    // 遍历分区表 4 个分区表项
    while(part_idx++ < 4) {
        if(p -> fs_type == 0x5) { // 扩展分区
            if(ext_lba_base != 0) { // 走逻辑分区
                partition_scan(hd, p -> start_lba + ext_lba_base);
            } else { // 走主分区
                ext_lba_base = p -> start_lba; // 记录扩展分区的地址
                partition_scan(hd, p -> start_lba);
            }
        } else if(p -> fs_type != 0) { // 合法的分区类型
            if(ext_lba == 0) { // 全为主分区
                hd -> prim_parts[p_no].start_lba = ext_lba + p -> start_lba;
                hd -> prim_parts[p_no].sec_cnt = p -> sec_cnt;
                hd -> prim_parts[p_no].my_disk = hd;
                list_append(&partition_list, &hd -> prim_parts[p_no].part_tag);
                sprintf(hd->prim_parts[p_no].name, "%s%d", hd->name, p_no + 1);
                p_no++;
                ASSERT(p_no < 4);
            } else { // 逻辑分区
                hd -> logic_parts[l_no].start_lba = ext_lba + p -> start_lba;
                hd -> logic_parts[l_no].sec_cnt = p -> sec_cnt;
                hd -> logic_parts[l_no].my_disk = hd;
                list_append(&partition_list, &hd -> logic_parts[l_no].part_tag);
                sprintf(hd->logic_parts[l_no].name, "%s%d", hd->name, l_no + 5); // 逻辑分区数字是从5开始,主分区是1～4.
                l_no++;
                if(l_no >= 8) // 我们只支持 8 个逻辑分区
                    return;
            }
        }
        p++;
    }
    sys_free(bs);
}

// 打印分区信息
static bool partition_info(struct list_elem* pelem, int arg UNUSED) {
   struct partition *part = elem2entry(struct partition, part_tag, pelem);
   printk("   %s start_lba:0x%x, sec_cnt:0x%x\n", part->name, part->start_lba, part->sec_cnt);

   /* 在此处return false与函数本身功能无关,
    * 只是为了让主调函数list_traversal继续向下遍历元素 */
   return false;

}

// 硬盘中断处理程序
void intr_hd_handler(uint8_t irq_no) {
    ASSERT(irq_no == 0x2e || irq_no == 0x2f);
    uint8_t ch_no = irq_no - 0x2e;
    struct ide_channel* channel = &channels[ch_no];
    ASSERT(channel -> irq_no == irq_no);

    /**若硬盘发生中断，则对应的接受者，即通道一定是唯一确认的
     * 因为为了避免无法分清中断信号来着同一通道上的哪块硬盘，咱们在操作硬盘时申请了锁
     */
    // 只需要硬盘中断，因此只要 expecting_intr != true，则不放行
    if(channel -> expecting_intr) {
        channel -> expecting_intr = false;
        sema_up(&channel -> disk_done); // 唤醒
        // 显示通知硬盘控制器本次中断已处理完成，不然硬盘不会继续产生新的中断
        inb(reg_status(channel)); // 这只是其中一种方法
    }
}

// 初始化硬盘数据结构
void ide_init() {
    printk("ide_init start.\n");
    uint8_t hd_cnt = *((uint8_t*)(0x475)); // 获取硬盘数量
    ASSERT(hd_cnt > 0);
    list_init(&partition_list);
    // 通道数量
    channel_cnt = DIV_ROUND_UP(hd_cnt, 2); // 一个 IDE 通道上有两个硬盘

    struct ide_channel* channel;
    uint8_t channel_no = 0, dev_no = 0;

    // 处理每个通道上的硬盘
    while(channel_no < channel_cnt) {
        channel = &channels[channel_no];
        sprintf(channel -> name, "ide%d", channel_no);

        // 为每个 IDE 通道初始化端口基地址及中断向量
        switch(channel_no) {
            case 0: // IDE0
                channel -> port_base = 0x1F0;  // 起始端口号的基数
                channel -> irq_no = 0x20 + 14; // 从片8259A上倒数第二个中断引脚，温盘，也就是 IDE0 通道上的中断向量号
                break;
            case 1: // IDE1
                channel -> port_base = 0x170;  // 起始端口号的基数
                channel -> irq_no = 0x20 + 15; // 从片8259A上最后一个中断引脚，我们用来响应 IDE1 通道上的硬盘中断
                break;
        }

        channel -> expecting_intr = false; // 未向硬盘发送命令，硬盘此时未工作，因此无需等待来自硬盘的中断
        lock_init(&channel -> lock);

        /* 初始化为0,目的是向硬盘控制器请求数据后,硬盘驱动sema_down此信号量会阻塞线程,
        直到硬盘完成后通过发中断,由中断处理程序将此信号量sema_up,唤醒线程. */
        sema_init(&channel -> disk_done, 0);

        register_handler(channel -> irq_no, intr_hd_handler);

        // 分别获取两个硬盘的参数及分区信息
        while(dev_no < 2) {
            struct disk* hd = &channel -> devices[dev_no];
            hd -> my_channel = channel;
            hd -> dev_no = dev_no;
            sprintf(hd->name, "sd%c", 'a' + channel_no * 2 + dev_no);
            identify_disk(hd); // 获取硬盘参数信息
            if(dev_no != 0) // 不对内核本身的裸硬盘(hd60M.img)进行处理
                partition_scan(hd, 0); // 扫描该硬盘上的所有分区
            p_no = 0, l_no = 0;
            dev_no++;
        }
        dev_no = 0;
        channel_no++; // 下一个 channel
    }

    printk("\n   all partition info\n");

    // 打印所有分区信息
    list_traversal(&partition_list, partition_info, (int) NULL);

    printk("ide_init done.\n");
}
